package client;

import common.*;

import java.io.IOException;
import java.net.InetSocketAddress;
import java.nio.ByteBuffer;
import java.nio.channels.SelectionKey;
import java.nio.channels.Selector;
import java.nio.channels.SocketChannel;
import java.util.Iterator;
import java.util.concurrent.*;
import java.util.logging.Level;
import java.util.logging.Logger;

/**
 * 客户端主类
 */
public class TunnelClient {
    private static final Logger logger = LogUtil.getLogger(TunnelClient.class.getName());

    // 配置
    private final ClientConfig config;

    // 本地服务连接管理
    private final ServiceManager serviceManager;

    // 控制通道
    private SocketChannel controlChannel;

    // 选择器
    private Selector selector;

    // 任务执行器
    private final Executor executor;

    // 定时任务执行器
    private final ScheduledExecutorService scheduler;

    // 心跳任务
    private ScheduledFuture<?> heartbeatTask;

    // 连接状态
    private boolean connected = false;

    // 认证状态
    private boolean authenticated = false;

    // 重试次数
    private int retryCount = 0;

    // 运行标志
    private volatile boolean running = false;

    private CountDownLatch latch = new CountDownLatch(1);

    /**
     * 构造函数
     */
    public TunnelClient(ClientConfig config) {
        this.config = config;
        this.serviceManager = new ServiceManager(this);
        this.executor = ConcurrencyUtil.createVirtualThreadExecutor();
        this.scheduler = ConcurrencyUtil.createScheduler(2, "client-scheduler");
    }

    /**
     * 启动客户端
     */
    public void start() throws InterruptedException {
        if (running) {
            logger.warning("Client is already running");
            return;
        }

        running = true;

        // 连接到服务器
        connect();

        logger.info("Tunnel client started");

        latch.await();
    }

    /**
     * 连接到服务器
     */
    private void connect() {
        // 在线程中执行连接操作
        executor.execute(() -> {
            try {
                // 创建控制通道
                controlChannel = SocketChannel.open();
                controlChannel.configureBlocking(false);

                // 设置TCP选项
                controlChannel.socket().setTcpNoDelay(true);
                controlChannel.socket().setKeepAlive(true);
                controlChannel.socket().setSoTimeout(config.getConnectionTimeout());

                // 连接到服务器
                logger.info("Connecting to server: " + config.getServerHost() + ":" + config.getServerPort());
                controlChannel.connect(new InetSocketAddress(config.getServerHost(), config.getServerPort()));

                // 初始化选择器
                selector = Selector.open();
                controlChannel.register(selector, SelectionKey.OP_CONNECT);

                // 处理连接过程
                processConnection();

            } catch (IOException e) {
                logger.log(Level.SEVERE, "Error connecting to server", e);
                scheduleReconnect();
            }
        });
    }

    /**
     * 处理连接过程
     */
    private void processConnection() {
        try {
            while (running && !connected) {
                selector.select(1000);

                Iterator<SelectionKey> keys = selector.selectedKeys().iterator();
                while (keys.hasNext()) {
                    SelectionKey key = keys.next();
                    keys.remove();

                    if (!key.isValid()) {
                        continue;
                    }

                    if (key.isConnectable()) {
                        finishConnect();
                    }
                }
            }

            // 连接成功后进行认证
            if (connected) {
                authenticate();
            }

        } catch (IOException e) {
            logger.log(Level.SEVERE, "Error in connection process", e);
            closeConnection();
            scheduleReconnect();
        }
    }

    /**
     * 完成连接过程
     */
    private void finishConnect() throws IOException {
        if (controlChannel.finishConnect()) {
            logger.info("Connected to server");

            // 更新通道注册的事件
            controlChannel.register(selector, SelectionKey.OP_READ);

            // 设置连接状态
            connected = true;
            retryCount = 0;
        }
    }

    /**
     * 进行认证
     */
    private void authenticate() {
        try {
            if (config.isAuthEnabled()) {
                logger.info("Authenticating with server");

                // 发送认证消息
                Message authMsg = Message.auth(config.getAuthToken());
                NetworkUtil.writeMessage(controlChannel, authMsg);

                // 等待认证响应
                Message response = NetworkUtil.readMessage(controlChannel);

                if (response.getType() == Protocol.MSG_TYPE_AUTH_RESP) {
                    byte result = response.getData()[0];
                    authenticated = (result == 1);

                    if (authenticated) {
                        logger.info("Authentication successful");
                        startMessageLoop();
                    } else {
                        logger.severe("Authentication failed");
                        closeConnection();
                        // 认证失败后不再重试
                        running = false;
                    }
                } else {
                    logger.warning("Unexpected response for authentication");
                    closeConnection();
                    scheduleReconnect();
                }
            } else {
                // 无需认证
                authenticated = true;
                startMessageLoop();
            }
        } catch (IOException e) {
            logger.log(Level.SEVERE, "Error during authentication", e);
            closeConnection();
            scheduleReconnect();
        }
    }

    /**
     * 开始消息处理循环
     */
    private void startMessageLoop() {
        // 启动心跳任务
        startHeartbeat();

        // 启动服务
        startServices();

        // 处理消息
        executor.execute(this::messageLoop);
    }

    /**
     * 启动心跳任务
     */
    private void startHeartbeat() {
        heartbeatTask = scheduler.scheduleAtFixedRate(() -> {
            try {
                if (connected && authenticated) {
                    // 发送心跳消息
                    NetworkUtil.writeMessage(controlChannel, Message.heartbeat());
                }
            } catch (IOException e) {
                logger.log(Level.WARNING, "Error sending heartbeat", e);
                closeConnection();
                scheduleReconnect();
            }
        }, 0, Protocol.HEARTBEAT_INTERVAL, TimeUnit.MILLISECONDS);
    }

    /**
     * 启动本地服务
     */
    private void startServices() {
        ServiceMapping[] mappings = config.getServiceMappings();
        if (mappings != null) {
            for (ServiceMapping mapping : mappings) {
                if (mapping != null) {
                    try {
                        serviceManager.addService(mapping);
                    } catch (IOException e) {
                        logger.log(Level.SEVERE, "Error starting service: " + mapping.getName(), e);
                    }
                }
            }
        }
    }

    /**
     * 消息处理循环
     */
    private void messageLoop() {
        try {
            while (running && connected && authenticated) {
                // 读取消息
                Message message = NetworkUtil.readMessage(controlChannel);

                // 处理消息
                handleMessage(message);
            }
        } catch (IOException e) {
            logger.log(Level.WARNING, "Error in message loop", e);
            closeConnection();
            scheduleReconnect();
        }
    }

    /**
     * 处理消息
     */
    private void handleMessage(Message message) {
        try {
            switch (message.getType()) {
                case Protocol.MSG_TYPE_HEARTBEAT:
                    // 心跳响应不需要特殊处理
                    break;
                case Protocol.MSG_TYPE_NEW_CONN:
                    handleNewConnection(message);
                    break;
                case Protocol.MSG_TYPE_DATA:
                    handleData(message);
                    break;
                case Protocol.MSG_TYPE_CLOSE_CONN:
                    handleCloseConnection(message);
                    break;
                default:
                    logger.warning("Unknown message type: " + message.getType());
            }
        } catch (Exception e) {
            logger.log(Level.SEVERE, "Error handling message", e);
        }
    }

    /**
     * 处理新连接消息
     */
    private void handleNewConnection(Message message) throws IOException {
        ByteBuffer buffer = ByteBuffer.wrap(message.getData());
        int connectionId = buffer.getInt();

        logger.info("New connection request: " + connectionId);

        // 创建本地服务连接
        serviceManager.createConnection(connectionId);
    }

    /**
     * 处理数据消息
     */
    private void handleData(Message message) throws IOException {
        ByteBuffer buffer = ByteBuffer.wrap(message.getData());
        int connectionId = buffer.getInt();

        // 提取数据部分
        byte[] data = new byte[buffer.remaining()];
        buffer.get(data);

        // 转发数据到本地服务
        serviceManager.forwardToLocalService(connectionId, data);
    }

    /**
     * 处理关闭连接消息
     */
    private void handleCloseConnection(Message message) {
        ByteBuffer buffer = ByteBuffer.wrap(message.getData());
        int connectionId = buffer.getInt();

        // 关闭本地服务连接
        serviceManager.closeConnection(connectionId);
    }

    /**
     * 发送数据到服务器
     */
    public void sendData(int connectionId, byte[] data) throws IOException {
        if (!connected || !authenticated) {
            throw new IOException("Not connected or authenticated");
        }

        Message dataMsg = Message.data(connectionId, data);
        NetworkUtil.writeMessage(controlChannel, dataMsg);
    }

    /**
     * 关闭连接
     */
    private void closeConnection() {
        if (!connected) {
            return;
        }

        connected = false;
        authenticated = false;

        // 停止心跳任务
        if (heartbeatTask != null) {
            heartbeatTask.cancel(false);
            heartbeatTask = null;
        }

        // 关闭所有本地服务连接
        serviceManager.closeAllConnections();

        try {
            if (controlChannel != null && controlChannel.isOpen()) {
                controlChannel.close();
            }

            if (selector != null) {
                selector.close();
            }

            logger.info("Disconnected from server");
        } catch (IOException e) {
            logger.log(Level.WARNING, "Error closing connection", e);
        }
    }

    /**
     * 安排重新连接
     */
    private void scheduleReconnect() {
        if (!running) {
            return;
        }

        retryCount++;

        if (config.getMaxRetries() > 0 && retryCount > config.getMaxRetries()) {
            logger.severe("Maximum retry count reached, stopping client");
            stop();
            return;
        }

        long delay = Protocol.RECONNECT_DELAY;

        logger.info("Scheduling reconnect attempt " + retryCount + " in " + (delay / 1000) + " seconds");

        scheduler.schedule(() -> {
            if (running) {
                connect();
            }
        }, delay, TimeUnit.MILLISECONDS);
    }

    /**
     * 停止客户端
     */
    public void stop() {
        if (!running) {
            return;
        }

        running = false;
        try {

            // 关闭连接
            closeConnection();

            // 关闭服务
            serviceManager.closeAllServices();

            // 关闭调度器
            ConcurrencyUtil.shutdownExecutor(scheduler);

            logger.info("Tunnel client stopped");
        } finally {
            latch.countDown();
        }

    }

    /**
     * 获取配置
     */
    public ClientConfig getConfig() {
        return config;
    }

    /**
     * 获取执行器
     */
    public Executor getExecutor() {
        return executor;
    }

    /**
     * 主方法
     */
    public static void main(String[] args) {
        try {
            // 加载配置
            String configPath = args.length > 0 ? args[0] : Thread.currentThread().getContextClassLoader().getResource("config/client.properties").getPath();
            Config configLoader = new Config();
            configLoader.load(configPath);

            // 创建客户端配置
            ClientConfig clientConfig = ClientConfig.fromConfig(configLoader);

            // 创建并启动客户端
            TunnelClient client = new TunnelClient(clientConfig);
            client.start();

            // 添加关闭钩子
            Runtime.getRuntime().addShutdownHook(new Thread(() -> {
                try {
                    client.stop();
                } catch (Exception e) {
                    e.printStackTrace();
                }
            }));

        } catch (Exception e) {
            System.err.println("Error starting client: " + e.getMessage());
            e.printStackTrace();
        }
    }
}
